package com.github.soiff.shadow.loader;

import com.github.soiff.shadow.utils.IOUtils;

import java.io.*;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.net.URLStreamHandlerFactory;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.regex.Pattern;

/**
 * Created by Soiff on 2016/05/20.
 *
 * @author : Soiff
 */
public class ResourceLoader extends URLClassLoader {

    final static Pattern any = Pattern.compile(".*");

    public ResourceLoader(URL[] urls, ClassLoader parent) {
        super(urls, parent);
    }

    public ResourceLoader(URL[] urls) {
        super(urls);
    }

    public ResourceLoader(URL[] urls, ClassLoader parent, URLStreamHandlerFactory factory) {
        super(urls, parent, factory);
    }

    public ResourceLoader(final File file, final File wd) throws IOException {
        super(load(file, wd));
    }

    public void addURLs(URL... urls) {
        for (int i = 0; null != urls && i < urls.length; ++ i)
            super.addURL(urls[i]);
    }

    public void addURLs(Collection<URL> urls) {
        if (null == urls) return;
        final Iterator<URL> it = urls.iterator();
        while (it.hasNext())
            super.addURL(it.next());
    }

    /**
     * Create a new ClassLoader from jar directly.
     * @param file
     * @param wd
     * @return
     * @throws MalformedURLException
     */
    private static URL[] load(File file, final File wd) throws IOException {
        /* Make sure working directory does not exist. */
        if (wd.exists())
            wd.delete();
        /* Create working directory and make sure it exists. */
        mkdir (wd);
        /* Uncompress jar to working directory. */
        uncompress (file, wd, any);
        /* We just need put several directories to working directory. */
        final List<URL> items = new ArrayList<>(32);
        items.add(new File(wd+"/").toURI().toURL());
        /* This would find single resource twice in loader.getResource(). */
        /* items.add(file.toURI().toURL()); */
        items.add(new File(wd, "classes/").toURI().toURL());
        File[] libs = new File(wd, "lib").listFiles();
        if (libs != null) {
            for (int i = 0; i < libs.length; i++) {
                items.add(libs[i].toURI().toURL());
            }
        }
        // create a normal parent-delegating classloader
        return items.toArray(new URL[0]);
    }

    /**
     * Unpack matching files from a jar. Entries inside the jar that do
     * not match the given pattern will be skipped.
     *
     * @param jarFile the .jar file to unpack
     * @param toDir the destination directory into which to unpack the jar
     * @param unpackRegex the pattern to match jar entries against
     */
    static List<URL> uncompress(File jarFile, File toDir, Pattern unpackRegex) throws IOException {
        final List<URL> items = new ArrayList<URL>(32);
        JarFile jar = new JarFile(jarFile);
        try {
            Enumeration<JarEntry> entries = jar.entries();
            while (entries.hasMoreElements()) {
                final JarEntry entry = entries.nextElement();
                if (!entry.isDirectory() && unpackRegex.matcher(entry.getName()).matches()) {
                    InputStream in = jar.getInputStream(entry);
                    try {
                        File file = new File(toDir, entry.getName());
                        mkdir(file.getParentFile());
                        OutputStream out = new FileOutputStream(file);
                        items.add(file.toURI().toURL());
                        try {
                            IOUtils.copyBytes(in, out, 8192);
                        } finally {
                            out.close();
                        }
                    } finally {
                        in.close();
                    }
                }
            }
        } finally {
            jar.close();
        }
        return items;
    }

    /**
     * Ensure the existence of a given directory.
     *
     * @throws IOException if it cannot be created and does not already exist
     */
    private static void mkdir(File dir) throws IOException {
        if (!dir.mkdirs() && !dir.isDirectory()) {
            throw new IOException("Mkdirs failed to create " + dir.toString());
        }
    }
}
